package com.fanshuai.tree.quadtree;

import java.util.ArrayList;
import java.util.List;

public class QuadTree implements IQuadTree {
    private static int NORTH_EAST = 1;
    private static int NORTH_WEST = 2;
    private static int SOUTH_EAST = 3;
    private static int SOUTH_WEST = 4;

    private static boolean isNorth(int dir) {
        return dir == NORTH_EAST || dir == NORTH_WEST;
    }

    class QuadNode {
        /**
         * 四叉树节点区域
         */
        private Rect region;

        /**
         * 区域内的点
         */
        private long[] X, Y;
        /**
         * 区域内的点数量
         */
        private int ptCount;

        /**
         * 四叉树节点的 西北、东北、西南、东南四个子节点
         */
        private QuadNode nw, ne, sw, se;

        public QuadNode(Rect rect) {
            if (null == rect) {
                throw new NullPointerException("rect is null");
            }

            this.region = rect;
            X = new long[POINT_NUM];
            Y = new long[POINT_NUM];
            ptCount = 0;
        }
    }

    /**
     * 默认单个区域最多保存的点数量
     */
    private int DEFAULT_POINT = 16;
    /**
     * 参数传入的单个区域最多保存的点数量
     */
    private int POINT_NUM = 16;

    private QuadNode root;

    public QuadTree(Rect r) {
        this.POINT_NUM = DEFAULT_POINT;
        root = new QuadNode(r);
    }

    /**
     * 构造函数
     * @param pointCount  每个区域存储的点数量
     */
    public QuadTree(Rect r, int pointCount) {
        this.POINT_NUM = pointCount;
        root = new QuadNode(r);
    }

    @Override
    public void add(long x, long y) {
        add(root, x, y);
    }

    /**
     * 当前节点表示的区域内添加(x,y)。如果当前节点已满，开辟子节点保存
     * @param x
     * @param y
     */
    private boolean add(QuadNode node, long x, long y) {
        Rect region = node.region;
        if (!region.contains(x, y)) {
            return false;
        }

        if (node.ptCount <= POINT_NUM) { //当前节点的point数组未满，直接保存
            node.X[node.ptCount] = x;
            node.Y[node.ptCount] = y;
            node.ptCount++;
            return true;
        } else { //开辟四个区域的子节点，并计算坐标，决定在其中某个子节点保存s
            //矩形中心坐标
            long midx = (region.x1 + region.x2) / 2;
            long midy = (region.y1 + region.y2) / 2;

            //开辟四个区域的子节点
            Rect rectNW = new Rect(region.x1, region.y1, midx, midy);
            Rect rectNE = new Rect(midx, region.y1, region.x2, midy);
            Rect rectSW = new Rect(region.x1, midy, midx, region.y2);
            Rect rectSE = new Rect(midx, midy, region.x2, region.y2);
            node.nw = new QuadNode(rectNW);
            node.ne = new QuadNode(rectNE);
            node.sw = new QuadNode(rectSW);
            node.se = new QuadNode(rectSE);

            //依次按照NW, NE, SW, SE顺序查询。(x,y)同时位于多个区域时(四个区域相交的直线)，按顺序只加入一个区域即可
            if (rectNW.contains(x, y)) {   //NW区域
                return add(node.nw, x, y);
            } else if (rectNE.contains(x, y)) {  //NE区域
                return add(node.ne, x, y);
            } else if (rectSW.contains(x, y)) {  //SW区域
                return add(node.sw, x, y);
            } else if (rectSE.contains(x, y)) {  //SE区域
                return add(node.se, x, y);
            }
        }
        return false;
    }

    @Override
    public int count(Rect rect) {
        return count(root, rect);
    }

    /**
     * 计算rect区域内节点数量
     * @param node
     * @param rect
     * @return
     */
    private int count(QuadNode node, Rect rect) {
        //矩形区域不相交，返回0
        if (!rect.intersect(node.region)) {
            return 0;
        }

        int count = 0;
        if (rect.contains(node.region)) {//查询区域完全包含当前节点的矩形区域
            count += node.ptCount;
        } else {
            for (int i = 0; i < node.ptCount; i++) {
                if (rect.contains(node.X[i], node.Y[i])) {
                    count++;
                }
            }
        }

        //计算四个子区域的点数量
        if (node.nw != null) {
            count += count(node.nw, rect);
            count += count(node.ne, rect);
            count += count(node.sw, rect);
            count += count(node.se, rect);
        }

        return count;
    }

    @Override
    public List<Point> points(Rect rect) {
        List<Point> points = new ArrayList<>();
        fillPoints(root, rect, points);
        return points;
    }

    /**
     * 获取矩形rect内所有的点
     * @param node
     * @param rect
     * @param points
     */
    private void fillPoints(QuadNode node, Rect rect, List<Point> points) {
        if (!rect.intersect(node.region)) {
            return;
        }

        if (rect.contains(node.region)) {
            for (int i = 0; i < node.ptCount; i++) {
                points.add(new Point(node.X[i], node.Y[i]));
            }
        } else {
            for (int i = 0; i < node.ptCount; i++) {
                if (rect.contains(node.X[i], node.Y[i])) {
                    points.add(new Point(node.X[i], node.Y[i]));
                }
            }
        }

        if (node.nw != null) {
            fillPoints(node.nw, rect, points);
            fillPoints(node.ne, rect, points);
            fillPoints(node.sw, rect, points);
            fillPoints(node.se, rect, points);
        }
    }

    @Override
    public List<Point> points(long x, long y, double radius) {
        Rect rect = new Rect((long)(x - radius), (long) (y - radius), (long)(x + radius), (long) (y + radius));
        List<Point> points = points(rect);

        Point center = new Point(x, y);

        //过滤
        List<Point> newPoints = new ArrayList<>();
        for (Point point : points) {
            if (Double.compare(point.distance(center), radius) < 0) {
                newPoints.add(point);
            }
        }
        return newPoints;
    }

    @Override
    public List<Point> getPoints() {
        List<Point> points = new ArrayList<>();
        fillPoints(root, points);
        return points;
    }

    /**
     * 获取所有的点
     * @param node
     * @param points
     */
    private void fillPoints(QuadNode node, List<Point> points) {
        if (null == node) {
            return;
        }

        for (int i = 0; i < node.ptCount; i++) {
            points.add(new Point(node.X[i], node.Y[i]));
        }
        if (node.nw != null) {
            fillPoints(node.nw, points);
            fillPoints(node.ne, points);
            fillPoints(node.sw, points);
            fillPoints(node.se, points);
        }
    }

    /**
     * 获取距离(x,y)最近的k个点
     * @param k
     * @param x
     * @param y
     * @return
     */
    @Override
    public List<Point> kNearestNeighbors(int k, long x, long y) {
        return null;
    }
}
